import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 或 'Heiti SC', 'Hiragino Sans GB'
plt.rcParams['axes.unicode_minus'] = False

# 1. 优化消融实验
df_ablation = pd.read_csv('./lesson3_v1/ablation_results.csv', encoding='utf-8')
# plt.figure(figsize=(10,5))
# sns.barplot(x='optimization', y='total_time_ms', data=df_ablation)
# plt.xticks(rotation=30, ha='right')
# plt.ylabel('总时间 (ms)')
# plt.title('不同优化策略下的推理+训练总时间')
# plt.tight_layout()
# plt.savefig('./lesson3_v1/figure/ablation_total_time.png')
# plt.close()

# plt.figure(figsize=(10,5))
# sns.barplot(x='optimization', y='speedup', data=df_ablation)
# plt.xticks(rotation=30, ha='right')
# plt.ylabel('加速比')
# plt.title('不同优化策略的加速比')
# plt.tight_layout()
# plt.savefig('./lesson3_v1/figure/ablation_speedup.png')
# plt.close()
plt.figure(figsize=(8, 6))
bars = plt.barh(df_ablation["optimization"], df_ablation["total_time_ms"], color='lightgreen')
plt.bar_label(bars, fmt='%.2f ms', padding=3)
plt.title('不同优化组合的总时间 (推理+训练)')
plt.xlabel('总时间 (ms)')
plt.grid(True, linestyle='--', alpha=0.6)
plt.tight_layout()
plt.savefig('./lesson3_v1/figure/ablation_total_time.png', dpi=300)
plt.close()


plt.figure(figsize=(8, 6))
bars = plt.barh(df_ablation["optimization"], df_ablation["speedup"], color='skyblue')
plt.bar_label(bars, fmt='%.2fx', padding=3)
plt.title('不同优化组合的加速比 (相对于基准版本)')
plt.xlabel('加速比')
plt.grid(True, linestyle='--', alpha=0.6)
plt.tight_layout()
plt.savefig('./lesson3_v1/figure/ablation_speedup.png', dpi=300)
plt.close()

# 2. 批大小实验
df_batch = pd.read_csv('./lesson3_v1/batch_size_results.csv', encoding='utf-8')
plt.figure(figsize=(8,5))
plt.plot(df_batch['batch_size'], df_batch['samples_per_second'], marker='o')
plt.xlabel('Batch Size')
plt.ylabel('吞吐量（样本/秒）')
plt.title('不同Batch Size下的吞吐量')
plt.grid()
plt.tight_layout()
plt.savefig('./lesson3_v1/figure/batch_size_throughput.png')
plt.close()

plt.figure(figsize=(8,5))
plt.plot(df_batch['batch_size'], df_batch['forward_time_ms'], marker='o', label='Forward')
plt.plot(df_batch['batch_size'], df_batch['backward_time_ms'], marker='o', label='Backward')
plt.xlabel('Batch Size')
plt.ylabel('时间 (ms)')
plt.title('Batch Size对前/后向传播时间影响')
plt.legend()
plt.grid()
plt.tight_layout()
plt.savefig('./lesson3_v1/figure/batch_size_time.png')
plt.close()

# 3. 隐藏层维度实验
df_hidden = pd.read_csv('./lesson3_v1/hidden_dim_results.csv', encoding='utf-8')
fig, ax1 = plt.subplots(figsize=(8,5))
ax1.plot(df_hidden['hidden_dim'], df_hidden['mse'], marker='o', color='C0', label='MSE')
ax1.set_xlabel('隐藏层维度')
ax1.set_ylabel('MSE', color='C0')
ax2 = ax1.twinx()
ax2.plot(df_hidden['hidden_dim'], df_hidden['total_time_ms'], marker='^', color='C1', label='总时间')
ax2.set_ylabel('总时间 (ms)', color='C1')
fig.suptitle('隐藏层维度对误差与耗时影响')
fig.tight_layout()
plt.savefig('./lesson3_v1/figure/hidden_dim_mse_time.png')
plt.close()

# 4. 损失曲线
df_loss = pd.read_csv('./lesson3_v1/training_history.csv')
plt.figure(figsize=(8,5))
plt.plot(df_loss['Epoch'], df_loss['TrainLoss'], label='Train Loss')
plt.plot(df_loss['Epoch'], df_loss['TestLoss'], label='Test Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('训练与测试集损失曲线')
plt.legend()
plt.grid()
plt.tight_layout()
plt.savefig('./lesson3_v1/figure/loss_curve.png')
plt.close()